package com.example.oauth2.service;

import com.example.oauth2.entry.AccessToken;
import com.example.oauth2.entry.RefreshToken;
import com.example.oauth2.repository.AccessTokenRepository;
import com.example.oauth2.repository.RefreshTokenRepository;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.common.util.SerializationUtils;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.stereotype.Component;

import java.io.UnsupportedEncodingException;
import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 自定义TokenStore类
 */
@Component
public class CustomTokenStore implements TokenStore {

    private Logger logger = LogManager.getLogger(getClass());

    @Autowired
    private AccessTokenRepository accessTokenRepository;
    @Autowired
    private RefreshTokenRepository refreshTokenRepository;

    private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();

    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {

        return readAuthentication(token.getValue());
    }

    @Override
    public OAuth2Authentication readAuthentication(String token) {

        OAuth2Authentication authentication = null;

        Optional<AccessToken> opt = Optional.ofNullable(accessTokenRepository.findOAuthAccessTokenByTokenId(token));
        if (opt.isPresent()) {
            authentication = deserializeAuthentication(opt.get().getAuthentication().getBytes());
        }

        return authentication;
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication authentication) {

        String refreshToken = null;
        if (token.getRefreshToken() != null) {
            refreshToken = token.getRefreshToken().getValue();
        }

        if (readAccessToken(token.getValue()) != null) {
            removeAccessToken(token.getValue());
        }

        AccessToken accessToken = new AccessToken();
        accessToken.setTokenId(token.getValue());
        accessToken.setToken(new String(serializeAccessToken(token)));
        accessToken.setAuthenticationId(authenticationKeyGenerator.extractKey(authentication));
        accessToken.setUsername(authentication.isClientOnly() ? null : authentication.getName());
        accessToken.setClientId(authentication.getOAuth2Request().getClientId());
        accessToken.setAuthentication(new String(serializeAuthentication(authentication)));
        accessToken.setRefreshToken(extractTokenKey(refreshToken));

        accessTokenRepository.save(accessToken);
    }

    @Override
    public OAuth2AccessToken readAccessToken(String tokenValue) {

        OAuth2AccessToken accessToken = null;

        Optional<AccessToken> opt = Optional.ofNullable(accessTokenRepository.findOAuthAccessTokenByTokenId(tokenValue));
        if (opt.isPresent()) {
            accessToken = deserializeAccessToken(opt.get().getToken().getBytes());
        }

        return accessToken;
    }

    @Override
    public void removeAccessToken(OAuth2AccessToken token) {

        removeAccessToken(token.getValue());
    }

    public void removeAccessToken(String tokenValue) {

        accessTokenRepository.deleteAllByTokenId(tokenValue);
    }

    @Override
    public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {

        RefreshToken oAuthRefreshToken = new RefreshToken();
        oAuthRefreshToken.setTokenId(refreshToken.getValue());
        oAuthRefreshToken.setToken(new String(serializeRefreshToken(refreshToken)));
        oAuthRefreshToken.setAuthentication(new String(serializeAuthentication(authentication)));

        refreshTokenRepository.save(oAuthRefreshToken);
    }

    @Override
    public OAuth2RefreshToken readRefreshToken(String tokenValue) {

        OAuth2RefreshToken refreshToken = null;

        Optional<RefreshToken> opt = Optional.ofNullable(refreshTokenRepository.findOAuthRefreshTokenByTokenId(tokenValue));
        if (opt.isPresent()) {
            refreshToken = deserializeRefreshToken(opt.get().getToken().getBytes());
        }

        return refreshToken;
    }

    @Override
    public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {

        return readAuthenticationForRefreshToken(token.getValue());
    }

    public OAuth2Authentication readAuthenticationForRefreshToken(String tokenValue) {

        OAuth2Authentication authentication = null;

        Optional<RefreshToken> opt = Optional.ofNullable(refreshTokenRepository.findOAuthRefreshTokenByTokenId(tokenValue));
        if (opt.isPresent()) {
            authentication = deserializeAuthentication(opt.get().getAuthentication().getBytes());
        }

        return authentication;
    }

    @Override
    public void removeRefreshToken(OAuth2RefreshToken token) {

        removeRefreshToken(token.getValue());
    }

    public void removeRefreshToken(String tokenValue) {

        refreshTokenRepository.deleteOAuthRefreshTokenByTokenId(tokenValue);
    }

    @Override
    public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {

        removeAccessTokenUsingRefreshToken(refreshToken.getValue());
    }

    public void removeAccessTokenUsingRefreshToken(String tokenValue) {

        accessTokenRepository.deleteAllByRefreshToken(extractTokenKey(tokenValue));
    }

    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {

        OAuth2AccessToken accessToken = null;

        String authenticationId = authenticationKeyGenerator.extractKey(authentication);
        Optional<AccessToken> opt = Optional.ofNullable(accessTokenRepository.findOAuthAccessTokenByAuthenticationId(authenticationId));
        if (opt.isPresent()) {
            accessToken = deserializeAccessToken(opt.get().getToken().getBytes());
            if (accessToken != null
                    && !authenticationId.equals(authenticationKeyGenerator.extractKey(readAuthentication(accessToken.getValue())))) {
                removeAccessToken(accessToken.getValue());
                // Keep the store consistent (maybe the same user is represented by this authentication but the details have
                // changed)
                storeAccessToken(accessToken, authentication);
            }
        }

        return accessToken;
    }

    @Override
    public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {

        List<OAuth2AccessToken> accessTokens = new ArrayList<OAuth2AccessToken>();
        Optional<List<AccessToken>> opt = Optional.ofNullable(accessTokenRepository.findAllByClientIdAndUsername(clientId, userName));
        if (opt.isPresent()) {
            accessTokens = opt.get().stream().map(AccessToken::getToken).filter(t -> t != null).map(t -> deserializeAccessToken(t.getBytes())).collect(Collectors.toList());
        }
        return accessTokens;
    }

    @Override
    public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {

        List<OAuth2AccessToken> accessTokens = new ArrayList<OAuth2AccessToken>();
        Optional<List<AccessToken>> opt = Optional.ofNullable(accessTokenRepository.findAllByClientId(clientId));
        if (opt.isPresent()) {
            accessTokens = opt.get().stream().map(AccessToken::getToken).filter(t -> t != null).map(t -> deserializeAccessToken(t.getBytes())).collect(Collectors.toList());
        }
        return accessTokens;
    }

    protected byte[] serializeAccessToken(OAuth2AccessToken token) {
        return Base64.getEncoder().encode(SerializationUtils.serialize(token));
    }

    protected byte[] serializeRefreshToken(OAuth2RefreshToken token) {
        return Base64.getEncoder().encode(SerializationUtils.serialize(token));
    }

    protected byte[] serializeAuthentication(OAuth2Authentication authentication) {
        return Base64.getEncoder().encode(SerializationUtils.serialize(authentication));
    }

    protected OAuth2AccessToken deserializeAccessToken(byte[] token) {
        return SerializationUtils.deserialize(Base64.getDecoder().decode(token));
    }

    protected OAuth2RefreshToken deserializeRefreshToken(byte[] token) {
        return SerializationUtils.deserialize(Base64.getDecoder().decode(token));
    }

    protected OAuth2Authentication deserializeAuthentication(byte[] authentication) {
        return SerializationUtils.deserialize(Base64.getDecoder().decode(authentication));
    }

    protected String extractTokenKey(String value) {
        if (value == null) {
            return null;
        }
        MessageDigest digest;
        try {
            digest = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException("MD5 algorithm not available.  Fatal (should be in the JDK).");
        }

        try {
            byte[] bytes = digest.digest(value.getBytes("UTF-8"));
            return String.format("%032x", new BigInteger(1, bytes));
        } catch (UnsupportedEncodingException e) {
            throw new IllegalStateException("UTF-8 encoding not available.  Fatal (should be in the JDK).");
        }
    }
}
